
#include <openrave-core.h>
#include <vector>
#include <iostream>

#include "orexample.h"
#include "pidcontroller.hpp"
#include "computedtorque.hpp"
#include "jointspacesolver.hpp"
#include "trajectory.hpp"
#include "spline.hpp"
#include "robotStruct.hpp"

#define TIMESTEP 10 //milliseconds

// Comment out to use simple PID,
// Leave uncommented to use unit mass compensation torque control: 
//#define ADVANCED_CONTROL 1

// Uncomment one only to select shape:
//#define WOBBLY_CIRCLE
//#define STRAIGHT_LINE
#define SPLINE_4POINT
//#define SPLINE_WANDER

// Instead of link1 having height, offset the origin
// This is to make simpler input
#define ZOFFSET 1.10

#define LINKLEN1 0.000
#define LINKLEN2 0.100
#define LINKLEN3 0.600
#define LINKLEN4 0.500
#define LINKLEN5 0.500

#define LINKWIDTH1 0.08
#define LINKWIDTH2 0.08
#define LINKWIDTH3 0.08
#define LINKWIDTH4 0.08     
#define LINKWIDTH5 0.08

#define LINKHEIGHT1 0.08
#define LINKHEIGHT2 0.08
#define LINKHEIGHT3 0.08
#define LINKHEIGHT4 0.08
#define LINKHEIGHT5 0.08

#define LINKCOM1 0.00
#define LINKCOM2 0.00
#define LINKCOM3 0.10
#define LINKCOM4 0.16
#define LINKCOM5 0.14

#define LINKMASS1 0.5
#define LINKMASS2 0.5
#define LINKMASS3 0.5
#define LINKMASS4 0.5
#define LINKMASS5 0.5

using std::cout;
using std::endl;
using std::string;

namespace cppexamples {

class RobotArm : public OpenRAVEExample {
 public:
  virtual void demothread(int argc, char ** argv) {
    
    // load the world:
    if( !penv->Load("models/world.xml") ){
      std::cout << "Failed to load xml\n";
      return;
    }
    // attach a physics engine:
    penv->SetPhysicsEngine(RaveCreatePhysicsEngine(penv, "ode"));
    penv->GetPhysicsEngine()->SetGravity(OpenRAVE::Vector(0,0,-9.8));
    
    // Find the first robot and get a pointer to it:
    std::vector<OpenRAVE::RobotBasePtr> orRobotVec;
    penv->GetRobots(orRobotVec);
    if(orRobotVec.size() < 1) {
      std::cout << "No robots found in xml\n";
      return;
    }
    
    OpenRAVE::RobotBasePtr orRobot = orRobotVec.at(0);
    uint32_t numJoints = orRobot->GetJoints().size();
    
    std::cout << "Init Robot Model:\n";
    RobotModel* robot = new RobotModel;
    robot->setLengths(LINKLEN1, LINKLEN2, LINKLEN3, LINKLEN4, LINKLEN5);
    robot->setMasses(LINKMASS1, LINKMASS2, LINKMASS3, LINKMASS4, LINKMASS5);
    robot->setWidths(LINKWIDTH1, LINKWIDTH2, LINKWIDTH3, LINKWIDTH4, LINKWIDTH5);
    robot->setHeights(LINKHEIGHT1, LINKHEIGHT2, LINKHEIGHT3, LINKHEIGHT4, LINKHEIGHT5);
    robot->setCOM(LINKCOM1, LINKCOM2, LINKCOM3, LINKCOM4, LINKCOM5);
    
    cout << "Forming Joint Trajectory...\n";
    assert(numJoints == 4);
    
    // ---------------- PLAN PATH ---------------- //
    
    JointSpaceSolver* solver = new JointSpaceSolver(robot);
    
    double reach = solver->getReachableLimit()-0.001;
    
    Trajectory* trajectory = new Trajectory();
    Setpoint* sp = new Setpoint(0, true);
    
    double x = reach/2;
    
    Spline* spline = new Spline();
    cvec4d* cvia = new cvec4d();
    
    // 4 POINT SPLINE
#ifdef SPLINE_4POINT
    cvia->set(reach, 0, 0, 0);
    spline->pushVia(cvia);
    cvia->set(x, -x, 0, M_PI/3);
    spline->pushVia(cvia);
    cvia->set(-x, -x, 0, M_PI/3);
    spline->pushVia(cvia);
    cvia->set(-x, x, 0, M_PI/3);
    spline->pushVia(cvia);
    cvia->set(reach/2, 0, 0, M_PI/2);
    spline->pushVia(cvia);
    
    spline->planPath(solver, trajectory, 4000, 20);
#endif //SPLINE_4POINT
    
    // WANDERING SPLINE
#ifdef SPLINE_WANDER
    cvia->set(reach, 0, 0, 0);
    spline->pushVia(cvia);
    cvia->set(x, 0, 0, M_PI/2);
    spline->pushVia(cvia);
    cvia->set(x/2, x, 0, M_PI/2);
    spline->pushVia(cvia);
    cvia->set(0, x/2, x/2, M_PI/3);
    spline->pushVia(cvia);
    cvia->set(0, x, -0.1, M_PI/3);
    spline->pushVia(cvia);
    cvia->set(-x, x, 0, M_PI/3);
    spline->pushVia(cvia);
    cvia->set(x, x/2, 0, M_PI/3);
    spline->pushVia(cvia);
    cvia->set(x, 0, -x/2, M_PI/2);
    spline->pushVia(cvia);
    cvia->set(reach, 0, 0, 0);
    spline->pushVia(cvia);
    
    spline->planPath(solver, trajectory, 4000, 20);
    
    printf("VIA POINTS:\n");
    spline->print();
#endif //SPLINE_WANDER
    
    delete spline;
    delete cvia;
    
    // ---- WOBBLY CIRCLE ----
#ifdef WOBBLY_CIRCLE
    sp->time = 0;
    sp->cpos->set(reach, 0.0, 0.0, M_PI/3);
    if(solver->solve(sp) != 0) return;
    trajectory->addSetpoint(sp);
    
    double path_r = 2*reach/3;
    double path_a = reach/12;
    
    sp->time = 1000;
    sp->cpos->set(path_r*cos(0), path_r*sin(0), 0, M_PI/3);
    if(solver->solve(sp) != 0) return;
    trajectory->addSetpoint(sp);
    
    for(double angle = 0; angle < 2*M_PI; angle += M_PI/180){
      sp->time = sp->time+100;
      double r = path_r + path_a*sin(angle*6);
      sp->cpos->set(r*cos(angle), r*sin(angle), 0, M_PI/3);
      if(solver->solve(sp) != 0) return;
      trajectory->addSetpoint(sp);
    }
#endif //WOBBLY_CIRCLE
    
    // ---- STRAIGHT LINE ----
#ifdef STRAIGHT_LINE
    double path_x = 3*reach/4;
    double path_y = reach/2;
    
    sp->time = 3000;
    sp->cpos->set(path_x, path_y, 0, M_PI/3);
    if(solver->solve(sp) != 0) return;
    trajectory->addSetpoint(sp);
    
    for(path_x = 3*reach/4; path_x > -3*reach/4; path_x -= 0.05){
      sp->time = sp->time+500;
      sp->cpos->set(path_x, path_y, 0, M_PI/3);
      if(solver->solve(sp) != 0) return;
      trajectory->addSetpoint(sp);
    }
#endif //STRAIGHT_LINE

    // Whole manipulator computed torque controller
#ifdef ADVANCED_CONTROL
    printf("Constructing Computed Torque Controller...\n");
    WholeManipController* controller = new WholeManipController(robot);

    // Set stiffness
    double stiffness = 10000.0;
    controller->setKp(stiffness);
    controller->setKv(2*sqrt(stiffness));     // Critical damping

#else // SIMPLE PID CONTROL

    // Construct Controllers:
    cout << "Constructing PID Controllers...\n";
    
    std::vector<PIDController*> controllers;
    controllers.resize(numJoints);
    
    // PID Gains:
    double Kp = 200.0;
    double Ki = 0.7;
    double Kd = 1000;
    
    // Max Joint Strength:
    double maxTorque = 100; // Nm
    
    // Joint 0:
    controllers[0] = new PIDController(Kp, Ki, Kd);
    controllers[0]->setAntiWindup(-maxTorque, maxTorque);

    // Joint 1:
    controllers[1] = new PIDController(Kp, Ki, Kd);
    controllers[1]->setAntiWindup(-maxTorque, maxTorque);
    
    // Joint 2:
    controllers[2] = new PIDController(Kp, Ki, Kd);
    controllers[2]->setAntiWindup(-maxTorque, maxTorque);
    
    // Joint 3:
    controllers[3] = new PIDController(Kp, Ki, Kd);
    controllers[3]->setAntiWindup(-maxTorque, maxTorque);
#endif

    cout << "Entering loop...\n";
    uint64_t simtime = 0.0;
    
    std::vector<OpenRAVE::dReal> jposVec;        // angular position
    jposVec.resize(numJoints);
    std::vector<OpenRAVE::dReal> jvelVec;       // angular velocities
    jvelVec.resize(numJoints);
    std::vector<OpenRAVE::dReal> torques;          // joint torques
    torques.resize(numJoints);
    
    // Loop until quit:
    while(IsOk()) {
      {  // Begin Mutually Exclusive Segment:
        OpenRAVE::EnvironmentMutex::scoped_lock lock(penv->GetMutex()); // lock environment
        
        simtime = penv->GetSimulationTime()/1000;    // get time and convert microseconds to milliseconds
        
        // Get Joint positions:
        orRobot->GetDOFValues(jposVec);
        orRobot->GetDOFVelocities(jvelVec);

        // Get Joint set points:
        trajectory->getSetpoint(simtime, sp);   // to repeat use (simtime%12000)
        
        // Do control:
#ifdef ADVANCED_CONTROL

        // Get control signal from advanced controller:
        controller->compute(jposVec, jvelVec, sp, torques);

#else // PID_CONTROL

        for(uint32_t i=0; i<numJoints; i++) {
          // Get control signal from PID:
          double controlSignal = controllers.at(i)->compute(jposVec.at(i), sp->jpos->get(i));
          
          // Max/Min to Actuator rails:
          double torque = controlSignal;
          if( controlSignal >  maxTorque ) torque =  maxTorque;
          if( controlSignal < -maxTorque ) torque = -maxTorque;
          torques.at(i) = torque;
        }

#endif // ADVANCED vs PID CONTROL

        // Set joint torques
        orRobot->SetDOFTorques(torques, false);
      }  // end Mutually Exclusive Segment
      
      // Do Drawing:
      std::list<OpenRAVE::GraphHandlePtr> listgraphs;
      listgraphs.clear();
      std::vector< OpenRAVE::RaveVector<float> > vpoints;
      
      OpenRAVE::RaveVector<float> white(1.0, 1.0, 1.0, 1.0);
      OpenRAVE::RaveVector<float> red(1.0, 0.0, 0.0, 1.0);
      OpenRAVE::RaveVector<float> black(0.0, 0.0, 0.0, 1.0);
      
      Setpoint* sp1 = new Setpoint(0, true);
      
      int dt = 400;
      
      //for(uint64_t ti = std::max((int)0, (int)simtime-10000); ti < simtime; ti += dt) {
      for(uint64_t ti = 0; ti < 40000; ti += dt) {
        vpoints.clear();
        
        trajectory->getSetpoint(ti, sp1);
        OpenRAVE::RaveVector<float> rv1(sp1->cpos->x, sp1->cpos->y, sp1->cpos->z+ZOFFSET);
        vpoints.push_back(rv1);
        
        trajectory->getSetpoint(ti + dt, sp1);
        OpenRAVE::RaveVector<float> rv2(sp1->cpos->x, sp1->cpos->y, sp1->cpos->z+ZOFFSET);
        vpoints.push_back(rv2);
        
        listgraphs.push_back(penv->drawlinestrip(&vpoints[0].x, vpoints.size(), sizeof(vpoints[0]), 1.0f, red));
      }
      delete sp1;
      vpoints.clear();
      
      OpenRAVE::RaveVector<float> rv3(0, 0, ZOFFSET);
      vpoints.push_back(rv3);
      
      cvec4d* cpos = new cvec4d();
      solver->solveFK(sp->jpos, cpos);
      
      OpenRAVE::RaveVector<float> rv4(cpos->x, cpos->y, cpos->z+ZOFFSET);
      vpoints.push_back(rv4);
      
      listgraphs.push_back(penv->drawlinestrip(&vpoints[0].x, vpoints.size(), sizeof(vpoints[0]), 1.0f, black));
      
      vpoints.clear();
      vpoints.push_back(rv3);
      
      jvec4d* jpos = new jvec4d(jposVec[0], jposVec[1], jposVec[2], jposVec[3]);
      
      solver->solveFK(jpos, cpos);
      OpenRAVE::RaveVector<float> rv6(cpos->x, cpos->y, cpos->z+ZOFFSET);
      vpoints.push_back(rv6);
      
      listgraphs.push_back(penv->drawlinestrip(&vpoints[0].x, vpoints.size(), sizeof(vpoints[0]), 1.0f, white));
      
      delete cpos;
      delete jpos;
      
      while( penv->GetSimulationTime()/1000 - simtime < TIMESTEP && IsOk()) {
        // wait...
      }
    }
    // Cleanup:
    delete sp;
    delete robot;
    delete solver;
    delete trajectory;

#ifdef ADVANCED_CONTROL
    delete controller;
#else // PID_CONTROL
    for(uint32_t i = 0; i < controllers.size(); i ++) {
      delete controllers.at(i);
    }
#endif

  }
};  // end class

} // end namespace cppexamples
